Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

flashinfer_bench/compile/registry.py (1)

55-64: Builder ordering LGTM; add a tiny trace to see which backend was chosen at runtime.

Priority Python > Triton > TVM‑FFI > CUDA is sensible. Consider logging the chosen builder once in BuilderRegistry.build for observability during benchmarks.

flashinfer_bench/compile/builders/__init__.py (1)

4-6: Sort all to satisfy Ruff RUF022.

Apply isort-style ordering.

-from .tvm_ffi_builder import TVMFFIBuilder
-
-__all__ = ["CUDABuilder", "PythonBuilder", "TritonBuilder", "TVMFFIBuilder"]
+from .tvm_ffi_builder import TVMFFIBuilder
+
+__all__ = ["CUDABuilder", "PythonBuilder", "TritonBuilder", "TVMFFIBuilder"]

Note: if you want strict lint compliance, change to alphabetical order:

-__all__ = ["CUDABuilder", "PythonBuilder", "TritonBuilder", "TVMFFIBuilder"]
+__all__ = ["CUDABuilder", "PythonBuilder", "TVMFFIBuilder", "TritonBuilder"]

As per static analysis hints.

flashinfer_bench/compile/__init__.py (1)

7-7: New exports are fine; keep all sorted to appease Ruff.

Sort entries alphabetically to pass RUF022.

-__all__ = [
-    "Builder",
-    "BuildError",
-    "BuilderRegistry",
-    "Runnable",
-    "get_builder_registry",
-    "PrebuiltLibraryManager",
-    "get_prebuilt_manager",
-]
+__all__ = [
+    "Builder",
+    "BuilderRegistry",
+    "BuildError",
+    "PrebuiltLibraryManager",
+    "Runnable",
+    "get_builder_registry",
+    "get_prebuilt_manager",
+]

examples/tvm_ffi_example.py (2)

84-101: JAX section likely runs on CPU arrays; ensure GPU device or skip.

For a CUDA kernel, place arrays on a GPU device before calling, or skip if no GPU is present to avoid runtime errors.

-        import jax.numpy as jnp
+        import jax, jax.numpy as jnp
@@
-        a_jax = jnp.array(a_torch.cpu().numpy())
-        b_jax = jnp.array(b_torch.cpu().numpy())
-        c_jax = jnp.empty((n,), dtype=jnp.float32)
+        gpus = jax.devices("gpu")
+        if not gpus:
+            raise ImportError("No JAX GPU device")
+        dev = gpus[0]
+        a_jax = jax.device_put(jnp.array(a_torch.cpu().numpy()), dev)
+        b_jax = jax.device_put(jnp.array(b_torch.cpu().numpy()), dev)
+        c_jax = jax.device_put(jnp.empty((n,), dtype=jnp.float32), dev)

---

9-13: Remove unused import.

flashinfer_bench as fib is unused.

-import flashinfer_bench as fib
 from flashinfer_bench.compile import get_builder_registry

flashinfer_bench/compile/prebuilt.py (3)

15-21: Docstring search order outdated.

The implementation also supports extra_paths and appends the local cache as step 4. Update the docstring to reflect the actual order.

---

103-110: Edge case: ensure get_cache_dir always returns the actual cache path even if it appeared earlier.

If an identical cache path is already in search_paths (e.g., via env/extra_paths), the “last element” assumption may break. Compute and return the canonical cache path directly.

-        cache_dir = self._search_paths[-1]  # Always the local cache
-        os.makedirs(cache_dir, exist_ok=True)
-        return cache_dir
+        cache_dir = os.path.join(get_fib_cache_path(), "tvm_ffi")
+        os.makedirs(cache_dir, exist_ok=True)
+        return cache_dir

---

88-101: Cross‑platform: only searches for .so.

Support .dylib (macOS) and .dll (Windows) to make prebuilt discovery portable.

-        # Try both with and without .so extension
-        if not lib_name.endswith(".so"):
-            filename = f"{lib_name}.so"
-        else:
-            filename = lib_name
-
-        for search_path in self._search_paths:
-            lib_path = os.path.join(search_path, filename)
-            if os.path.exists(lib_path):
-                logger.debug(f"Found prebuilt library: {lib_path}")
-                return lib_path
+        # Try common extensions across platforms
+        candidates = [lib_name] if any(lib_name.endswith(ext) for ext in (".so", ".dylib", ".dll")) \
+            else [f"{lib_name}{ext}" for ext in (".so", ".dylib", ".dll")]
+        for search_path in self._search_paths:
+            for fname in candidates:
+                lib_path = os.path.join(search_path, fname)
+                if os.path.exists(lib_path):
+                    logger.debug(f"Found prebuilt library: {lib_path}")
+                    return lib_path

tests/compile/test_prebuilt.py (1)

79-93: Make tests portable across OS library suffixes.

Hardcoding .so breaks on Windows/macOS. Parameterize the extension or probe via multiple suffixes.

-            lib_path = os.path.join(tmpdir, "test_lib.so")
+            import sys
+            ext = ".dll" if sys.platform.startswith("win") else (".dylib" if sys.platform == "darwin" else ".so")
+            lib_path = os.path.join(tmpdir, f"test_lib{ext}")
@@
-            found2 = manager.find("test_lib.so")
+            found2 = manager.find(f"test_lib{ext}")

tests/compile/test_tvm_ffi_builder.py (2)

178-179: Escape regex metacharacters in match=.

Use a raw string or escape dots so the message matches literally.

-with pytest.raises(BuildError, match="No .cu CUDA sources"):
+with pytest.raises(BuildError, match=r"No \.cu CUDA sources"):

As per static analysis hint (RUF043).

---

114-122: Avoid relying on private method in test (builder._get_lib_path).

Accessing a private method couples tests to internals. Prefer asserting on the returned path from build (e.g., via runnable.meta["cache_dir"] and expected filename) or expose a small public helper.

-        lib_path = builder._get_lib_path(simple_solution)
+        cache_dir = builder._prebuilt_manager.get_cache_dir()
+        # Derive expected filename from solution name if needed, or assert any *.so exists under cache_dir.
+        lib_path = next((p for p in (os.path.join(cache_dir, f) for f in os.listdir(cache_dir)) if p.endswith(".so")), None)

flashinfer_bench/compile/builders/tvm_ffi_builder.py (2)

49-64: Use unpacking for list concatenation.

On line 64, prefer unpacking over concatenation for better performance and readability.

Apply this diff:

-                elif sys.platform == "win32":
-                    ldflags = [f"/LIBPATH:{lib_path}"] + lib_names
+                elif sys.platform == "win32":
+                    ldflags = [f"/LIBPATH:{lib_path}", *lib_names]

---

66-69: Consider logging exceptions during package path discovery.

The bare except Exception: pass silently swallows all errors during package path discovery, making it difficult to diagnose issues when dependencies are misconfigured or packages are malformed.

Apply this diff:

     except Exception:
-        pass
+        logger.debug(f"Failed to discover package paths for {pkg_name}")
 
     return include_path, ldflags

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

flashinfer_bench/compile/builder.py (1)

33-45: Hash suffix length change is reasonable; note cache key impact

Bumping the SHA-1 suffix from 4 to 6 hex chars improves collision resistance for package names and should be harmless, but it will stop reusing any artifacts keyed with the old suffix length (they’ll just be rebuilt and old cache entries remain orphaned on disk). If that cache is user-visible or disk-sensitive, consider a one-time cleanup or a short note in release docs.

flashinfer_bench/compile/runnable.py (3)

10-20: Runnable meta is now stored by reference instead of copied

Storing meta directly instead of copying it means any external mutations to the dict (or sharing the same dict across multiple runnables) will be visible here. If meta is treated as immutable metadata this is fine, but if callers may mutate it, consider documenting the expectation or reinstating a defensive copy (dict(meta or {})) to avoid surprising aliasing.

---

75-82: Unify close() behavior with base Runnable

TVMFFIRunnable.close() currently reimplements the close logic without the try/finally guard used in Runnable.close(). While functionally similar in the non-error case, using super().close() would keep the behavior consistent (e.g., ensuring _closer is nulled even if it raises):

-    def close(self) -> None:
-        if self._closer:
-            self._closer()
-            self._closer = None
+    def close(self) -> None:
+        super().close()

This also centralizes any future changes to close semantics in the base class.

---

41-73: Confirmed: torch.empty device optimization is valid for torch 2.8.0

The dest-passing logic is correct and well-structured:

• Uses Definition.get_var_values on input tensor shapes, then get_output_shapes to size outputs
• Allocates torch.empty with appropriate dtype via dtype_str_to_torch_dtype and places outputs on the first input's device
• Returns single tensor vs. list based on output count, mirroring base Runnable behavior

Two clarifications worth documenting:

• All input kwargs are assumed to be torch tensors with .shape and .device attributes (non-tensor or mixed-device inputs would misbehave)
• Outputs are always allocated on the first input's device; if heterogeneous devices or explicit selection appears later, consider adding a device parameter or validation

PyTorch 2.8.0 supports torch.empty with a device parameter that allocates the tensor directly on the specified device, so optimize the allocation:

-            output_tensors[name] = torch.empty(
-                shape, dtype=dtype_str_to_torch_dtype(self._definition.outputs[name].dtype)
-            ).to(device)
+            output_tensors[name] = torch.empty(
+                shape,
+                dtype=dtype_str_to_torch_dtype(self._definition.outputs[name].dtype),
+                device=device,
+            )

flashinfer_bench/compile/builders/__init__.py (1)

4-6: LGTM: TVMFFIBuilder integration.

The import and export are correct and properly integrate the new TVM-FFI builder.

Note: Static analysis suggests sorting __all__ alphabetically for consistency. This is purely stylistic and can be deferred.

-__all__ = ["CUDABuilder", "PythonBuilder", "TritonBuilder", "TVMFFIBuilder"]
+__all__ = ["CUDABuilder", "PythonBuilder", "TVMFFIBuilder", "TritonBuilder"]

flashinfer_bench/compile/builders/tvm_ffi_builder.py (1)

74-76: Optional: Simplify unused parameter.

The cpp_files parameter is unused in _get_language. While this maintains API symmetry, consider simplifying if the parameter won't be needed.

-    def _get_language(self, cpp_files: List[str], cuda_files: List[str]) -> str:
+    def _get_language(self, cuda_files: List[str]) -> str:
         return "cuda" if len(cuda_files) > 0 else "cpp"

And update the call site at line 112.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

tests/compile/test_tvm_ffi_builder.py (1)

94-94: Remove redundant import.

torch is already imported at line 7. This redundant import should be removed.

Apply this diff:

-    import torch
-

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

flashinfer_bench/compile/builders/tvm_ffi_builder.py (2)

22-34: Fix docstring inconsistency with actual implementation.

The docstring on line 27 mentions tvm_ffi.cpp.build_inline(), but the implementation on line 117 actually uses tvm_ffi.cpp.build(). This inconsistency can confuse maintainers.

Update the docstring:

     Build strategy:
     1. Check if .so exists in cache (multi-process safe)
-    2. If not, compile with tvm_ffi.cpp.build_inline() to cache
+    2. If not, compile with tvm_ffi.cpp.build() to cache
     3. Load with tvm_ffi.load_module()

---

115-123: Remove or complete the incomplete comment.

Line 116 contains an incomplete comment: "Use build_inline instead of build to". This should either be completed with the rationale or removed to avoid confusion.

Apply this fix:

         try:
-            # Use build_inline instead of build to
+            # Build the TVM-FFI module with caching
             output_lib_path = tvm_ffi.cpp.build(

flashinfer_bench/data/definition.py (1)

114-115: Consider adding validation for non-overlapping output/input names.

The docstring now states that output names must not overlap with input names, but there's no corresponding validation in the model validators to enforce this constraint. If this is a hard requirement, consider adding a validator to prevent violations at definition-creation time.

Add a validator if this constraint should be enforced:

     @model_validator(mode="after")
     def _validate_tensor_axis_references(self) -> "Definition":
         """Validate that tensor shapes reference defined axes.
 
         Ensures that all axis names used in input and output tensor shapes
         are properly defined in the axes dictionary.
 
         Raises
         ------
         ValueError
             If any tensor shape references an undefined axis.
         """
         all_tensors = {**self.inputs, **self.outputs}
 
         for tensor_name, tensor_spec in all_tensors.items():
             if tensor_spec.shape is not None:
                 for axis_name in tensor_spec.shape:
                     if axis_name not in self.axes:
                         tensor_type = "input" if tensor_name in self.inputs else "output"
                         raise ValueError(
                             f'{tensor_type.capitalize()} "{tensor_name}" references undefined '
                             f'axis "{axis_name}"'
                         )
         return self
+
+    @model_validator(mode="after")
+    def _validate_no_overlapping_names(self) -> "Definition":
+        """Validate that output names do not overlap with input names."""
+        overlapping = set(self.inputs.keys()) & set(self.outputs.keys())
+        if overlapping:
+            raise ValueError(
+                f"Output names must not overlap with input names. Found: {overlapping}"
+            )
+        return self

flashinfer_bench/compile/builders/__init__.py (1)

1-6: Consider sorting all alphabetically.

The static analysis tool suggests sorting __all__ for consistency. While not critical, alphabetical ordering improves maintainability.

Apply this diff:

-__all__ = ["CUDABuilder", "PythonBuilder", "TritonBuilder", "TVMFFIBuilder"]
+__all__ = ["CUDABuilder", "PythonBuilder", "TVMFFIBuilder", "TritonBuilder"]

flashinfer_bench/compile/builders/tvm_ffi_builder.py (2)

41-42: Consider adding TVM-FFI availability check.

Similar to CUDABuilder.can_build() which checks CUDA availability, TVMFFIBuilder.can_build() should verify that TVM-FFI is available in the environment. Currently, it only checks the language, which means compilation failures will occur at build time rather than being caught earlier by can_build().

Consider adding an availability check pattern similar to CUDABuilder:

 class TVMFFIBuilder(Builder):
+    _tvm_ffi_available: bool = None
+
+    @classmethod
+    def _get_tvm_ffi_available(cls) -> bool:
+        if cls._tvm_ffi_available is None:
+            try:
+                import tvm_ffi
+                import tvm_ffi.cpp
+                cls._tvm_ffi_available = True
+            except ImportError:
+                cls._tvm_ffi_available = False
+        return cls._tvm_ffi_available
+
     def can_build(self, sol: Solution) -> bool:
-        return sol.spec.language == SupportedLanguages.CUDA
+        return sol.spec.language == SupportedLanguages.CUDA and self._get_tvm_ffi_available()

---

74-75: Address unused parameter flagged by static analysis.

The static analysis tool flags cpp_files as unused in _get_language(). While the current logic (return "cuda" if cuda_files else "cpp") is correct, the parameter is indeed unused. Consider whether this is intentional for API consistency or if the logic should be updated.

If the parameter should remain for API consistency, add a comment or underscore prefix. Otherwise, the logic could be:

-def _get_language(self, cpp_files: List[str], cuda_files: List[str]) -> str:
-    return "cuda" if len(cuda_files) > 0 else "cpp"
+def _get_language(self, cpp_files: List[str], cuda_files: List[str]) -> str:
+    if len(cuda_files) > 0:
+        return "cuda"
+    elif len(cpp_files) > 0:
+        return "cpp"
+    else:
+        # Should not reach here due to validation in _write_sources
+        return "cpp"

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro